import sys
sys.path.append("../")
from model.StockBaseInfo import *
import pandas as pd
import matplotlib.pyplot as plt
from mpl_finance import candlestick_ohlc
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import json
from matplotlib.dates import DateFormatter, MonthLocator
from frameworks.utils.RedisUtil import *

mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
import datetime


def get_data(code):
    redis = RedisUtil()
    stockbaseinfo = StockBaseInfo()
    # txt = "概念板块： " + stockbaseinfo.getGaiNian(code) + "<br>--------------------------------------------<br>"
    key = "dfcf_data_" + str(code)
    al = redis.vget(key)
    if al == None:
        print("=================set cache=======================")
        klineday = stockbaseinfo.getDfcfLineDay(code, 30)
        al = json.dumps(klineday)
        redis.vset(key, al)
    else:
        print("=======read cache===============")
    zf = 0
    dates = []
    m = 0
    klineday = json.loads(al)
    if "klines" in klineday:
        data = []
        for row in klineday["klines"]:
            daydata = row.split(",")
            if m == 0:
                dates.append(daydata[0])
            info = {
                "code": code,
                "trade_date": daydata[0],
                "open": daydata[1],
                "high": daydata[3],
                "low": daydata[4],
                "close": daydata[2],
                "volume": round(float(daydata[5]) / 10000, 2),
                "money": round(float(daydata[6]) / 100000000, 2),
                "zf": float(daydata[8]),
                "zfprice": daydata[9],
                "turn": daydata[10]
            }
            data.append(info)
        newdfdata = pd.DataFrame(data)
        ma_list = [5, 10, 20]
        for ma in ma_list:
            newdfdata['ma' + str(ma)] = newdfdata["close"].rolling(ma).mean()
        newdfdata.set_index("trade_date", inplace=True)
        df_sorted = newdfdata.sort_index(ascending=True)
        return df_sorted
    else:
        return False

def get_data_no_index(code, start_day, end_day, calc_day=""):
    redis = RedisUtil()
    stockbaseinfo = StockBaseInfo()
    # txt = "概念板块： " + stockbaseinfo.getGaiNian(code) + "<br>--------------------------------------------<br>"
    key = "dfcf_data_" + str(code)
    al = redis.vget(key)
    if al == None:
        print("=================set cache=======================")
        klineday = stockbaseinfo.getDfcfLineDay(code, 30)
        al = json.dumps(klineday)
        redis.vset(key, al)
    else:
        print("=======read cache===============")
    zf = 0
    dates = []
    m = 0
    klineday = json.loads(al)
    if "klines" in klineday:
        data = []
        for row in klineday["klines"]:
            daydata = row.split(",")
            if m == 0:
                dates.append(daydata[0])
            info = {
                "code": code,
                "trading_date": daydata[0],
                "open": float(daydata[1]),
                "high": float(daydata[3]),
                "low": float(daydata[4]),
                "close": float(daydata[2]),
                "volume": round(float(daydata[5]) / 10000, 2),
                "money": round(float(daydata[6]) / 100000000, 2),
                "zf": float(daydata[8]),
                "zfprice": daydata[9],
                "turn": daydata[10]
            }
            data.append(info)
        newdfdata = pd.DataFrame(data)
        newdfdata.set_index("trading_date", inplace=True)
        dfall = newdfdata.sort_index(ascending=True)
        dfall['Date'] = list(
            map(lambda x: mdates.date2num(datetime.datetime.strptime(x, '%Y-%m-%d')), dfall.index.tolist()))
        dfall['ma3'] = dfall["close"].rolling(3).mean()
        dfall['ma5'] = dfall["close"].rolling(5).mean()
        dfall['ma10'] = dfall["close"].rolling(10).mean()
        dfall['ma20'] = dfall["close"].rolling(20).mean()
        return dfall
    else:
        return False

def callBuyFunc(row,pre,three):
    #return float(row["close"]) >= float(row["ma5"]) and float(row["ma5"]) > float(row["ma10"]) and float(row["ma10"]) > float(row["ma18"]) and float(row["ma18"]) > float(row["ma30"])
    #return float(row["ma5"]) > float(row["ma10"]) and float(row["ma10"]) > float(row["ma18"]) and float(row["ma18"]) > float(row["ma30"])
    #return float(row["ma5"]) > float(row["ma10"]) and row["volume"] > pre["volume"] and pre["volume"] > three["volume"]
    return float(row["ma5"]) > float(pre["ma5"]) and (float(row["ma5"])-float(three["ma5"]))/float(three["ma5"]) > 0.008 and float(row["close"]) > float(row["ma5"])

def callSellFunc(row,pre,three):
    #return float(row["close"]) >= float(row["ma5"]) and float(row["ma5"]) > float(row["ma10"]) and float(row["ma10"]) > float(row["ma18"]) and float(row["ma18"]) > float(row["ma30"])
    #return float(row["ma5"]) > float(row["ma10"]) and float(row["ma10"]) > float(row["ma18"]) and float(row["ma18"]) > float(row["ma30"])
    #return float(row["ma5"]) > float(row["ma10"]) and row["volume"] > pre["volume"] and pre["volume"] > three["volume"]
    return float(row["ma5"]) < float(row["ma10"])

def onequery(code, argv):
    account = 100000
    hasmoney = 100000
    df = get_data(code)
    buy = {}
    buymoney = 0
    pre = {}
    three = {}
    i = 0
    lastcnt = 0
    for date, row in df.iterrows():
        if i == 0:
            pre = row
            i += 1
            continue
        if i == 1:
            three = pre
            pre = row
            i += 1
            continue
        if i < 10:
            three = pre
            pre = row
            i += 1
            continue
        if callBuyFunc(row,pre,three) and len(buy) == 0:
            if (float(row["close"]) <= 0):
                continue
            print("=====================买入===========================")
            print(float(row["ma5"]))
            print(float(pre["ma5"]))
            print(float(three["ma5"]))
            print((float(row["ma5"])-float(three["ma5"]))/float(three["ma5"]))
            buynum = int(float(account) / (float(row["close"]) * 100))
            buy = {"code": code, "holdnum": buynum, "account": account}
            hasmoney = float(account) - float(row["close"]) * int(buynum) * 100
            print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + "," + code + ",买入" + str(
                    buy["holdnum"] * 100) + ",持有现金" + str(hasmoney))
        else:
            if callBuyFunc(row,pre,three) == False and len(buy) > 0:
                account = float(buy["holdnum"]) * 100 * float(row["close"]) + float(hasmoney)
                hasmoney = account
                print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + "," + code + ",卖出" + str(
                        buy["holdnum"] * 100) + ",持有现金" + str(hasmoney))
                buy = {}
                #lastcnt = i + 5
            elif len(buy) > 0:
                account = float(buy["holdnum"]) * 100 * float(row["close"]) + float(hasmoney)
                print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + "," + code + ",持有股份" + str(
                        buy["holdnum"] * 100) + ",持有现金" + str(hasmoney))
                pass
            else:
                print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + ",持有现金" + str(hasmoney))
                pass
        three = pre
        pre = row
        i += 1
    if len(buy) > 0:
        print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + "," + code + ",持有" + str(buy["holdnum"] * 100) + ",持有现金" + str(
            hasmoney))
    else:
        print("日期:" + date + "，价格:" + str(row["close"]) + "，总账户:" + str(account) + "," + code + ",持有0,持有现金" + str(hasmoney))


def showImage(stock_code, data_price):
    # 5、绘制成交量
    fig = plt.figure(figsize=(18, 10))
    # plt.style.use("dark_background")
    grid = plt.GridSpec(12, 10, wspace=0.5, hspace=0.5)
    # （1）绘制K线图
    # K线数据
    ohlc = data_price[['Date', 'open', 'high', 'low', 'close']]
    ohlc.loc[:, 'Date'] = range(len(ohlc))  # 重新赋值横轴数据，绘制K线图无间隔
    # 绘制K线
    ax1 = fig.add_subplot(grid[0:8, 0:12])  # 设置K线图的尺寸
    candlestick_ohlc(ax1, ohlc.values.tolist(), width=.7
                     , colorup='red', colordown='green')
    ax1.plot(range(len(data_price)), data_price['ma5']
             , color='black', lw=2, label='MA (5)')
    ax1.plot(range(len(data_price)), data_price['ma10']
             , color='blue', lw=2, label='MA (10)')
    ax1.plot(range(len(data_price)), data_price['ma20']
             , color='green', lw=2, label='MA (20)')
    plt.title(stock_code, fontsize=14)  # 设置图片标题
    plt.ylabel('价 格（元）', fontsize=14)  # 设置纵轴标题
    plt.legend(loc='best')
    ax1.set_xticks([])  # 日期标注在成交量中，故清空此处x轴刻度
    ax1.set_xticklabels([])  # 日期标注在成交量中，故清空此处x轴
    ax1.grid(True)  ##设置网格线
    ax1.xaxis_date()
    # （2）绘制成交量
    # 成交量数据
    data_volume = data_price[['Date', 'close', 'open', 'volume']]
    data_volume['color'] = data_volume.apply(lambda row: 1 if row['close'] >= row['open'] else 0,
                                             axis=1)  # 计算成交量柱状图对应的颜色，使之与K线颜色一致
    data_volume.Date = ohlc.Date
    # 绘制成交量
    ax2 = fig.add_subplot(grid[8:10, 0:12])  # 设置成交量图形尺寸
    ax2.bar(data_volume.query('color==1')['Date']
            , data_volume.query('color==1')['volume']
            , color='r')  # 绘制红色柱状图
    ax2.bar(data_volume.query('color==0')['Date']
            , data_volume.query('color==0')['volume']
            , color='g')  # 绘制绿色柱状图
    plt.xticks(rotation=30)
    plt.xlabel('日 期', fontsize=14)  # 设置横轴标题
    # 修改横轴日期标注
    date_list = ohlc.index.tolist()  # 获取日期列表
    xticks_len = round(len(date_list) / (len(ax2.get_xticks()) - 1))  # 获取默认横轴标注的间隔
    xticks_num = range(0, len(date_list), xticks_len)  # 生成横轴标注位置列表
    xticks_str = list(map(lambda x: date_list[int(x)], xticks_num))  # 生成正在标注日期列表
    ax2.set_xticks(xticks_num)  # 设置横轴标注位置
    ax2.set_xticklabels(xticks_str)  # 设置横轴标注日期
    ax2.grid(True)
    plt.legend()
    plt.show()


def showImag2(code, str_day, end_day):
    newdf = get_data_no_index(code, str_day, end_day)
    newdf.index = pd.to_datetime(newdf.index)
    fig, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(15, 8))
    ax1.plot(newdf["close"], label='close')
    ax1.plot(newdf["ma20"], label='ma20')
    ax1.plot(newdf["ma10"], label='ma10')
    ax1.set_title('wandayuanxian' + code)
    ax1.set_ylabel('Price')
    ax1.grid(True)
    ax1.xaxis_date()
    ax2.plot(newdf["v20"], label='v20')
    ax2.plot(newdf["v10"], label='v10')
    ax2.plot(newdf["v5"], label='v5')
    ax2.bar(newdf.index, newdf["volume"], width=0.5)
    ax2.set_ylabel('Volume')
    ax2.grid(True)
    plt.legend()
    plt.show()
    # cla = PltImage()
    # cla.run(codes[i], orgargs["start_day"], orgargs["end_day"])

code = sys.argv[1]
onequery(code,{"start_day":"2021-01-01","end_day":"2021-12-01"})
data = get_data_no_index(code, "2021-01-01", "2021-12-01")
showImage(code, data)

